"""
数据库管理模块
"""
import asyncio
import json
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict

from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, Boolean, Text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.sql import func

from .config import settings

logger = logging.getLogger(__name__)

# 创建数据库引擎
engine = create_engine(
    settings.DATABASE_URL,
    echo=settings.DEBUG,
    pool_size=10,
    max_overflow=20,
    pool_timeout=30,
    pool_recycle=3600
)

# 创建会话工厂
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建基类
Base = declarative_base()


# 数据模型定义
class StrategyRecord(Base):
    """策略记录表"""
    __tablename__ = "strategies"
    
    id = Column(Integer, primary_key=True, index=True)
    strategy_id = Column(String(50), unique=True, index=True)
    name = Column(String(100))
    strategy_type = Column(String(50))
    parameters = Column(Text)  # JSON格式
    is_active = Column(Boolean, default=False)
    created_at = Column(DateTime, default=func.now())
    updated_at = Column(DateTime, default=func.now(), onupdate=func.now())


class MarketDataRecord(Base):
    """行情数据表"""
    __tablename__ = "market_data"
    
    id = Column(Integer, primary_key=True, index=True)
    symbol = Column(String(20), index=True)
    timestamp = Column(DateTime, index=True)
    open_price = Column(Float)
    high_price = Column(Float)
    low_price = Column(Float)
    close_price = Column(Float)
    volume = Column(Integer)
    turnover = Column(Float)
    bid_price = Column(Float)
    ask_price = Column(Float)
    bid_volume = Column(Integer)
    ask_volume = Column(Integer)
    delta = Column(Float)
    gamma = Column(Float)
    vega = Column(Float)
    theta = Column(Float)
    implied_volatility = Column(Float)


class TradeRecord(Base):
    """交易记录表"""
    __tablename__ = "trades"
    
    id = Column(Integer, primary_key=True, index=True)
    strategy_id = Column(String(50), index=True)
    order_id = Column(String(100), unique=True, index=True)
    symbol = Column(String(20), index=True)
    trade_type = Column(String(10))  # buy/sell
    quantity = Column(Integer)
    price = Column(Float)
    amount = Column(Float)
    commission = Column(Float)
    trade_time = Column(DateTime, index=True)
    status = Column(String(20), default="completed")  # pending/filled/cancelled/rejected
    created_at = Column(DateTime, default=func.now())


class PositionRecord(Base):
    """持仓记录表"""
    __tablename__ = "positions"
    
    id = Column(Integer, primary_key=True, index=True)
    strategy_id = Column(String(50), index=True)
    symbol = Column(String(20), index=True)
    quantity = Column(Integer)
    avg_price = Column(Float)
    current_price = Column(Float)
    market_value = Column(Float)
    pnl = Column(Float)
    pnl_rate = Column(Float)
    created_at = Column(DateTime, default=func.now())
    updated_at = Column(DateTime, default=func.now(), onupdate=func.now())


class RiskAlertRecord(Base):
    """风险告警记录表"""
    __tablename__ = "risk_alerts"
    
    id = Column(Integer, primary_key=True, index=True)
    alert_type = Column(String(50))
    severity = Column(String(20))  # low/medium/high/critical
    message = Column(Text)
    data = Column(Text)  # JSON格式
    created_at = Column(DateTime, default=func.now())


class BacktestRecord(Base):
    """回测记录表"""
    __tablename__ = "backtests"
    
    id = Column(Integer, primary_key=True, index=True)
    strategy_id = Column(String(50), index=True)
    start_date = Column(DateTime)
    end_date = Column(DateTime)
    initial_capital = Column(Float)
    final_capital = Column(Float)
    total_return = Column(Float)
    annual_return = Column(Float)
    max_drawdown = Column(Float)
    sharpe_ratio = Column(Float)
    win_rate = Column(Float)
    total_trades = Column(Integer)
    parameters = Column(Text)  # JSON格式
    results = Column(Text)  # JSON格式
    created_at = Column(DateTime, default=func.now())


@dataclass
class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self):
        self.engine = engine
        self.SessionLocal = SessionLocal
    
    async def initialize(self):
        """初始化数据库"""
        try:
            # 创建所有表
            Base.metadata.create_all(bind=self.engine)
            logger.info("数据库表创建完成")
            
            # 检查是否需要初始化数据
            await self._initialize_default_data()
            
        except Exception as e:
            logger.error(f"数据库初始化失败: {e}")
            raise
    
    async def _initialize_default_data(self):
        """初始化默认数据"""
        try:
            with self.SessionLocal() as session:
                # 检查是否已有策略记录
                existing_strategies = session.query(StrategyRecord).count()
                
                if existing_strategies == 0:
                    # 创建默认策略记录
                    default_strategies = [
                        StrategyRecord(
                            strategy_id="ma_cross_50etf",
                            name="50ETF均线交叉策略",
                            strategy_type="moving_average",
                            parameters=json.dumps({
                                "short_period": 5,
                                "long_period": 20,
                                "rsi_period": 14,
                                "rsi_overbought": 70,
                                "rsi_oversold": 30
                            }),
                            is_active=False
                        ),
                        StrategyRecord(
                            strategy_id="vol_breakout_300etf",
                            name="300ETF波动率突破策略",
                            strategy_type="volatility_breakout",
                            parameters=json.dumps({
                                "volatility_period": 20,
                                "breakout_multiplier": 1.5,
                                "stop_loss": 0.02,
                                "take_profit": 0.05
                            }),
                            is_active=False
                        )
                    ]
                    
                    session.add_all(default_strategies)
                    session.commit()
                    logger.info("默认策略数据初始化完成")
                
        except Exception as e:
            logger.error(f"初始化默认数据失败: {e}")
    
    async def close(self):
        """关闭数据库连接"""
        try:
            self.engine.dispose()
            logger.info("数据库连接已关闭")
        except Exception as e:
            logger.error(f"关闭数据库连接失败: {e}")
    
    def get_session(self) -> Session:
        """获取数据库会话"""
        return self.SessionLocal()
    
    async def save_strategy(self, strategy_data: Dict) -> bool:
        """保存策略"""
        try:
            with self.SessionLocal() as session:
                # 检查是否已存在
                existing = session.query(StrategyRecord).filter_by(
                    strategy_id=strategy_data["strategy_id"]
                ).first()
                
                if existing:
                    # 更新现有记录
                    existing.name = strategy_data["name"]
                    existing.strategy_type = strategy_data["strategy_type"]
                    existing.parameters = json.dumps(strategy_data.get("parameters", {}))
                    existing.is_active = strategy_data.get("is_active", False)
                    existing.updated_at = datetime.now()
                else:
                    # 创建新记录
                    strategy = StrategyRecord(
                        strategy_id=strategy_data["strategy_id"],
                        name=strategy_data["name"],
                        strategy_type=strategy_data["strategy_type"],
                        parameters=json.dumps(strategy_data.get("parameters", {})),
                        is_active=strategy_data.get("is_active", False)
                    )
                    session.add(strategy)
                
                session.commit()
                return True
                
        except Exception as e:
            logger.error(f"保存策略失败: {e}")
            return False
    
    async def get_strategy(self, strategy_id: str) -> Optional[Dict]:
        """获取策略"""
        try:
            with self.SessionLocal() as session:
                strategy = session.query(StrategyRecord).filter_by(
                    strategy_id=strategy_id
                ).first()
                
                if strategy:
                    return {
                        "strategy_id": strategy.strategy_id,
                        "name": strategy.name,
                        "strategy_type": strategy.strategy_type,
                        "parameters": json.loads(strategy.parameters),
                        "is_active": strategy.is_active,
                        "created_at": strategy.created_at,
                        "updated_at": strategy.updated_at
                    }
                return None
                
        except Exception as e:
            logger.error(f"获取策略失败: {e}")
            return None
    
    async def get_all_strategies(self) -> List[Dict]:
        """获取所有策略"""
        try:
            with self.SessionLocal() as session:
                strategies = session.query(StrategyRecord).all()
                
                return [{
                    "strategy_id": s.strategy_id,
                    "name": s.name,
                    "strategy_type": s.strategy_type,
                    "parameters": json.loads(s.parameters),
                    "is_active": s.is_active,
                    "created_at": s.created_at,
                    "updated_at": s.updated_at
                } for s in strategies]
                
        except Exception as e:
            logger.error(f"获取所有策略失败: {e}")
            return []
    
    async def save_market_data(self, data: Dict) -> bool:
        """保存行情数据"""
        try:
            with self.SessionLocal() as session:
                market_data = MarketDataRecord(
                    symbol=data["symbol"],
                    timestamp=data["timestamp"],
                    open_price=data.get("open", 0),
                    high_price=data.get("high", 0),
                    low_price=data.get("low", 0),
                    close_price=data.get("close", 0),
                    volume=data.get("volume", 0),
                    turnover=data.get("turnover", 0),
                    bid_price=data.get("bid_price", 0),
                    ask_price=data.get("ask_price", 0),
                    bid_volume=data.get("bid_volume", 0),
                    ask_volume=data.get("ask_volume", 0),
                    delta=data.get("delta", 0),
                    gamma=data.get("gamma", 0),
                    vega=data.get("vega", 0),
                    theta=data.get("theta", 0),
                    implied_volatility=data.get("implied_volatility", 0)
                )
                
                session.add(market_data)
                session.commit()
                return True
                
        except Exception as e:
            logger.error(f"保存行情数据失败: {e}")
            return False
    
    async def get_market_data(self, symbol: str, limit: int = 100) -> List[Dict]:
        """获取行情数据"""
        try:
            with self.SessionLocal() as session:
                data = session.query(MarketDataRecord).filter_by(
                    symbol=symbol
                ).order_by(MarketDataRecord.timestamp.desc()).limit(limit).all()
                
                return [{
                    "symbol": d.symbol,
                    "timestamp": d.timestamp,
                    "open": d.open_price,
                    "high": d.high_price,
                    "low": d.low_price,
                    "close": d.close_price,
                    "volume": d.volume,
                    "turnover": d.turnover,
                    "bid_price": d.bid_price,
                    "ask_price": d.ask_price,
                    "bid_volume": d.bid_volume,
                    "ask_volume": d.ask_volume,
                    "delta": d.delta,
                    "gamma": d.gamma,
                    "vega": d.vega,
                    "theta": d.theta,
                    "implied_volatility": d.implied_volatility
                } for d in data]
                
        except Exception as e:
            logger.error(f"获取行情数据失败: {e}")
            return []
    
    async def save_trade_record(self, trade_data: Dict) -> bool:
        """保存交易记录"""
        try:
            with self.SessionLocal() as session:
                trade = TradeRecord(
                    strategy_id=trade_data["strategy_id"],
                    order_id=trade_data["order_id"],
                    symbol=trade_data["symbol"],
                    trade_type=trade_data["trade_type"],
                    quantity=trade_data["quantity"],
                    price=trade_data["price"],
                    amount=trade_data["amount"],
                    commission=trade_data.get("commission", 0),
                    trade_time=trade_data["trade_time"],
                    status=trade_data.get("status", "completed")
                )
                
                session.add(trade)
                session.commit()
                return True
                
        except Exception as e:
            logger.error(f"保存交易记录失败: {e}")
            return False
    
    async def get_trade_records(self, strategy_id: Optional[str] = None, limit: int = 100) -> List[Dict]:
        """获取交易记录"""
        try:
            with self.SessionLocal() as session:
                query = session.query(TradeRecord)
                
                if strategy_id:
                    query = query.filter_by(strategy_id=strategy_id)
                
                trades = query.order_by(TradeRecord.trade_time.desc()).limit(limit).all()
                
                return [{
                    "id": t.id,
                    "strategy_id": t.strategy_id,
                    "order_id": t.order_id,
                    "symbol": t.symbol,
                    "trade_type": t.trade_type,
                    "quantity": t.quantity,
                    "price": t.price,
                    "amount": t.amount,
                    "commission": t.commission,
                    "trade_time": t.trade_time,
                    "status": t.status,
                    "created_at": t.created_at
                } for t in trades]
                
        except Exception as e:
            logger.error(f"获取交易记录失败: {e}")
            return []
    
    async def save_position(self, position_data: Dict) -> bool:
        """保存持仓"""
        try:
            with self.SessionLocal() as session:
                # 检查是否已存在
                existing = session.query(PositionRecord).filter_by(
                    strategy_id=position_data["strategy_id"],
                    symbol=position_data["symbol"]
                ).first()
                
                if existing:
                    # 更新现有记录
                    existing.quantity = position_data["quantity"]
                    existing.avg_price = position_data["avg_price"]
                    existing.current_price = position_data["current_price"]
                    existing.market_value = position_data["market_value"]
                    existing.pnl = position_data["pnl"]
                    existing.pnl_rate = position_data["pnl_rate"]
                    existing.updated_at = datetime.now()
                else:
                    # 创建新记录
                    position = PositionRecord(
                        strategy_id=position_data["strategy_id"],
                        symbol=position_data["symbol"],
                        quantity=position_data["quantity"],
                        avg_price=position_data["avg_price"],
                        current_price=position_data["current_price"],
                        market_value=position_data["market_value"],
                        pnl=position_data["pnl"],
                        pnl_rate=position_data["pnl_rate"]
                    )
                    session.add(position)
                
                session.commit()
                return True
                
        except Exception as e:
            logger.error(f"保存持仓失败: {e}")
            return False
    
    async def get_positions(self, strategy_id: Optional[str] = None) -> List[Dict]:
        """获取持仓"""
        try:
            with self.SessionLocal() as session:
                query = session.query(PositionRecord)
                
                if strategy_id:
                    query = query.filter_by(strategy_id=strategy_id)
                
                positions = query.all()
                
                return [{
                    "id": p.id,
                    "strategy_id": p.strategy_id,
                    "symbol": p.symbol,
                    "quantity": p.quantity,
                    "avg_price": p.avg_price,
                    "current_price": p.current_price,
                    "market_value": p.market_value,
                    "pnl": p.pnl,
                    "pnl_rate": p.pnl_rate,
                    "created_at": p.created_at,
                    "updated_at": p.updated_at
                } for p in positions]
                
        except Exception as e:
            logger.error(f"获取持仓失败: {e}")
            return []
    
    async def save_risk_alert(self, alert_data: Dict) -> bool:
        """保存风险告警"""
        try:
            with self.SessionLocal() as session:
                alert = RiskAlertRecord(
                    alert_type=alert_data["alert_type"],
                    severity=alert_data.get("severity", "medium"),
                    message=alert_data.get("message", ""),
                    data=json.dumps(alert_data.get("data", {}))
                )
                
                session.add(alert)
                session.commit()
                return True
                
        except Exception as e:
            logger.error(f"保存风险告警失败: {e}")
            return False
    
    async def get_risk_alerts(self, limit: int = 100) -> List[Dict]:
        """获取风险告警"""
        try:
            with self.SessionLocal() as session:
                alerts = session.query(RiskAlertRecord).order_by(
                    RiskAlertRecord.created_at.desc()
                ).limit(limit).all()
                
                return [{
                    "id": a.id,
                    "alert_type": a.alert_type,
                    "severity": a.severity,
                    "message": a.message,
                    "data": json.loads(a.data) if a.data else {},
                    "created_at": a.created_at
                } for a in alerts]
                
        except Exception as e:
            logger.error(f"获取风险告警失败: {e}")
            return []
    
    async def save_backtest_result(self, backtest_data: Dict) -> bool:
        """保存回测结果"""
        try:
            with self.SessionLocal() as session:
                backtest = BacktestRecord(
                    strategy_id=backtest_data["strategy_id"],
                    start_date=backtest_data["start_date"],
                    end_date=backtest_data["end_date"],
                    initial_capital=backtest_data["initial_capital"],
                    final_capital=backtest_data["final_capital"],
                    total_return=backtest_data["total_return"],
                    annual_return=backtest_data["annual_return"],
                    max_drawdown=backtest_data["max_drawdown"],
                    sharpe_ratio=backtest_data["sharpe_ratio"],
                    win_rate=backtest_data["win_rate"],
                    total_trades=backtest_data["total_trades"],
                    parameters=json.dumps(backtest_data.get("parameters", {})),
                    results=json.dumps(backtest_data.get("results", {}))
                )
                
                session.add(backtest)
                session.commit()
                return True
                
        except Exception as e:
            logger.error(f"保存回测结果失败: {e}")
            return False
    
    async def get_backtest_results(self, strategy_id: Optional[str] = None, limit: int = 10) -> List[Dict]:
        """获取回测结果"""
        try:
            with self.SessionLocal() as session:
                query = session.query(BacktestRecord)
                
                if strategy_id:
                    query = query.filter_by(strategy_id=strategy_id)
                
                backtests = query.order_by(BacktestRecord.created_at.desc()).limit(limit).all()
                
                return [{
                    "id": b.id,
                    "strategy_id": b.strategy_id,
                    "start_date": b.start_date,
                    "end_date": b.end_date,
                    "initial_capital": b.initial_capital,
                    "final_capital": b.final_capital,
                    "total_return": b.total_return,
                    "annual_return": b.annual_return,
                    "max_drawdown": b.max_drawdown,
                    "sharpe_ratio": b.sharpe_ratio,
                    "win_rate": b.win_rate,
                    "total_trades": b.total_trades,
                    "parameters": json.loads(b.parameters) if b.parameters else {},
                    "results": json.loads(b.results) if b.results else {},
                    "created_at": b.created_at
                } for b in backtests]
                
        except Exception as e:
            logger.error(f"获取回测结果失败: {e}")
            return []


# 全局数据库管理器实例
db_manager = DatabaseManager()
